// 2025-07-14

#include <iostream>
#include <string>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/dnn/dnn.hpp>
#include <opencv2/dnn/all_layers.hpp>
#include <onnxruntime_cxx_api.h>
#include <opencv2/core/utils/logger.hpp>
#include <cmath> // 包含 round 函数
#include <chrono>

// #define LOG true // 是否打印日志
#define LOG false // 是否打印日志
using namespace std;

size_t HOLE_NUM_W = 12; // 12;            // 穴盘列数
size_t HOLE_NUM_H = 6; // 6;            // 穴盘行数
float MATCH_THERSHOLD = 0.5f;     // 孔,苗匹配 相交阈值
float AREA_THERSHOLD = 0.15f;      // 判断弱苗阈值， 弱苗占穴孔的比例
float SCORE_THERSHOLD = 0.4f;     // 预测的置信度阈值
float IOU = 0.5f;                // NMS IOU 阈值

#define COLOR_WEAK cv::Scalar(0, 0, 255)   // 弱苗的颜色
#define COLOR_EXITS cv::Scalar(0, 255, 0)  // 正常苗的颜色



// 获取掩码和掩码面积
void drawBox(cv::Mat& img, std::vector<float>& b_xyxyc);
float iou(const std::vector<float> hole_box1, const std::vector<float> box2);
bool get_hole_xyxys(std::vector<std::vector<float>>& hole_xyxys, int w, int h);


/************************定义类 */
class model {
public:

    model(const std::string& onnxpath); //构造函数
    void preProcess(cv::Mat& img);
    std::vector<int> infer(cv::Mat img);
    void postProcess(std::vector<int>& res);
    void forward();
    void save_img(const std::string save_path);
    void hole_target_match(const std::vector<std::vector<float>>& target_xyxys, std::vector<bool>& holes_match_bool);

public:
    Ort::Env env; //  创建一个环境
    std::unique_ptr<Ort::Session> session; // 创建一个会话
    std::array<int64_t, 4> input_shape_info; // 输入尺度
    size_t input_size; // 输入大小
    // 创建 input_names 和 output_names， 输入输出名称
    std::vector<const char*> input_names;
    std::vector<const char*> output_names;
    std::vector<std::string> output_node_names;
    std::vector<std::string> input_node_names;
    int input_w, input_h, output_h, output_w, img_w, img_h; // 输入输出原始尺寸
    cv::Mat img_processed; // 输入图片
    cv::Mat mFrame; // 原始图片
    float m_x_ratio, m_y_ratio; // 缩放比例
    size_t mask_size;
    std::vector<std::vector<float>> hole_xyxys;
    float* pdata_res;
};




int main() {
    cv::utils::logging::setLogLevel(cv::utils::logging::LOG_LEVEL_ERROR); // 屏蔽opencv警告

    // 1 创建一个模型20250427limiao-seg_simplified_model
    model model("D:\\codes\\deploy_codes\\weights\\limiao-seg\\limiao-detect.onnx");
    // 2 图片文件路径
    //std::string filename = "D:\\codes\\deploy_codes\\datas\\limiao\\18.jpg";
    std::string filename = "D:\\codes\\deploy_codes\\datas\\limiao\\21.bmp";
    // 3. 读取文件
    cv::Mat img = cv::imread(filename);

    // 4. 推理
    auto start = std::chrono::high_resolution_clock::now();

    std::vector<int> res;
    res = model.infer(img);

    auto end = std::chrono::high_resolution_clock::now();
    // 计算耗时（毫秒）
    auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << "time: " << duration_ms << " ms" << std::endl;
    model.save_img("./res.jpg");  // 保存 结果图片


    // 5. 打印输出
    for (int i = 0; i < res.size(); i++) {
        if (i % HOLE_NUM_W == 0) {
            printf("\n");
        }
        printf("%d ", res[i]);
    }


    std::cout << "LOG test end !" << std::endl;


    return 0;
}




model::model(const std::string& onnxpath) {
    std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());

    Ort::SessionOptions session_options;  // 创建一个会话选项
    env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "inference"); // 创建一个环境
    session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL);  // 设置优化级别/、
    //session_options.SetIntraOpNumThreads(8);  // 设置内核线程数
    //session_options.SetInterOpNumThreads(4);  // 设置外部线程数
    //OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);

    session = std::make_unique<Ort::Session>(env, modelPath.c_str(), session_options); // 创建一个会话

    size_t numInputNodes = session->GetInputCount(); // 获取输入节点数量
    size_t numOutputNodes = session->GetOutputCount(); // 获取输出节点数量
    Ort::AllocatorWithDefaultOptions allocator; // 创建一个分配器

    // 获取输入节点名称和类型信息， 主要用来获取输入节点的名称
    const char* name_cstr;
    Ort::TypeInfo input_type_info = session->GetInputTypeInfo(0);
    auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
    auto input_dims = input_tensor_info.GetShape();
    input_w = input_dims[2];
    input_h = input_dims[3];
    for (int i = 0; i < numInputNodes; i++) {
        auto input_name = session->GetInputNameAllocated(i, allocator); // 第 i 个输入节点的名称，并由 allocator 管理其内存
        name_cstr = input_name.get();
        input_node_names.push_back(name_cstr);

    }

    // 获取输出节点名称和类型信息
    Ort::TypeInfo output_type_info = session->GetOutputTypeInfo(0);
    auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
    auto output_dims = output_tensor_info.GetShape();
    output_h = output_dims[1];  // 输出维度8400
    output_w = output_dims[2];  // 25605
    for (int i = 0; i < numOutputNodes; i++) {
        auto out_name = session->GetOutputNameAllocated(i, allocator);
        name_cstr = out_name.get();
        output_node_names.push_back(name_cstr);

    }


    input_names.resize(input_node_names.size());
    output_names.resize(output_node_names.size());
    for (size_t i = 0; i < input_node_names.size(); ++i) { input_names[i] = input_node_names[i].c_str(); }
    // 填充 output_names
    for (size_t i = 0; i < output_node_names.size(); ++i) { output_names[i] = output_node_names[i].c_str(); }

    // 输入信息构建
    input_size = input_h * input_w * 3;
    input_shape_info = { 1, 3, input_h, input_w };
    mask_size = input_w / 4; // mask 输出的大小



};

void model::preProcess(cv::Mat& img) {

    mFrame = img;
    img_w = img.cols;
    img_h = img.rows;
    m_x_ratio = static_cast<float>(img.cols) / input_w; // rate
    m_y_ratio = static_cast<float>(img.rows) / input_h;
    img_processed = cv::dnn::blobFromImage(img, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(), true, false, CV_32F);

    get_hole_xyxys(hole_xyxys, img_w, img_h);  // 获取洞洞位置


}

void model::forward() {

    auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);

    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator_info,
        img_processed.ptr<float>(),
        input_size, input_shape_info.data(),
        input_shape_info.size());


    std::vector<Ort::Value> ort_outputs;
    if (!session) {
        // 确保session_已经被初始化
        throw std::runtime_error("LOG Session is not initialized. Call Initialize first.");
    }

    try {
        ort_outputs = session->Run(Ort::RunOptions{ nullptr },
            input_names.data(),
            &input_tensor,
            input_names.size(),
            output_names.data(),
            output_names.size());
    }
    catch (std::exception e) {
        std::cout << e.what() << std::endl;
    }


    pdata_res = ort_outputs[0].GetTensorMutableData<float>();

};


void model::postProcess(std::vector<int>& res) {

    // 1. 获取输出数据
    size_t out_num = 8400;
    // 2. 解析结果
    std::vector<float> scores; // 置信度分数
    std::vector<cv::Rect> bboxes;
    std::vector<std::vector<float>> masks;

    // 预分配空间pdata_res
    size_t output_w = 25605;
    float* line;
    for (size_t i = 0; i < out_num; i++) {
        line = pdata_res + output_w * i; // 当前样本的起始指针

        if (*(line + 4) < SCORE_THERSHOLD) {
            continue;
        }

        // 1. 分数
        scores.push_back(*(line + 4)); // 预测分数
        // 2. rectangle
        cv::Rect bbox(*(line + 0), *(line + 1), *(line + 2), *(line + 3));
        bboxes.emplace_back(bbox);

        // 3. mask
        std::vector<float> mask_(line + 5, line + output_w);
        masks.emplace_back(mask_);
    }


    // 3. NMS
    std::vector<int> indices;  // 存储保留的边界框
    cv::dnn::NMSBoxes(bboxes, scores, SCORE_THERSHOLD, IOU, indices);




    // 4. 解析mask， 还原到原始图像
    std::vector<cv::Mat> res_masks;
    std::vector<std::vector<float>> res_bboxes;
    std::vector<float> res_areas;


    float hole_area = img_w / HOLE_NUM_W * img_h / HOLE_NUM_H;
    for (int idx : indices) {
        // 1. 处理box
        cv::Rect box = bboxes[idx];
        float x1, x2, y1, y2;
        float xc = (float)box.x;
        float yc = (float)box.y;
        float w = (float)box.width;
        float h = (float)box.height;

        // 1. box 
        x1 = (xc - w / 2) / input_w * img_w;
        y1 = (yc - h / 2) / input_h * img_h;
        x2 = (xc + w / 2) / input_w * img_w;
        y2 = (yc + h / 2) / input_h * img_h;
        std::vector<float> box_xyxy = { x1,y1,x2,y2 };


        // 2. 处理mask
        std::vector<float> mask_ = masks[idx];
        int mask_x1 = static_cast<int>(std::round(x1 / img_w * mask_size)); // 在mask上对应的坐标
        int mask_y1 = static_cast<int>(std::round(y1 / img_h * mask_size));
        int mask_x2 = static_cast<int>(std::round(x2 / img_w * mask_size));
        int mask_y2 = static_cast<int>(std::round(y2 / img_h * mask_size));
        cv::Mat mask(mask_size, mask_size, CV_32F, mask_.data()); // 160*160
        cv::Mat cropped_mask = mask(cv::Range(mask_y1, mask_y2), cv::Range(mask_x1, mask_x2));  // 根据160*160剪切目标mask
        int target_width = static_cast<int>(std::round(x2 - x1)); // 最终目标结果尺寸
        int target_height = static_cast<int>(std::round(y2 - y1));
        cv::resize(cropped_mask, cropped_mask, cv::Size(target_width, target_height), 0, 0, cv::INTER_LINEAR); // 还原实际图像大小比例
        // 构建mask 0-1 矩阵
        cv::Mat binary_mask;
        cv::threshold(cropped_mask, binary_mask, 0.5f, 255, cv::THRESH_BINARY);
        binary_mask.convertTo(binary_mask, CV_8U);


        // 3. 计算面积 筛选苗  画出苗的范围
        float mask_area = (float)cv::countNonZero(binary_mask);  // 非0值面积
        if ((mask_area / hole_area) > AREA_THERSHOLD) {
            res_bboxes.push_back(box_xyxy);
            res_masks.push_back(binary_mask);
            res_areas.push_back(mask_area);
        }
        else {
            // 4. 画图
            cv::Mat roi = mFrame(cv::Rect(x1, y1, x2 - x1, y2 - y1));
            //cv::Scalar color = (mask_area / hole_area) < AREA_THERSHOLD ? COLOR_WEAK : COLOR_EXITS;
            std::vector<std::vector<cv::Point>> contours;
            cv::Mat hierarchy;
            cv::findContours(binary_mask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE); // 查找轮廓
            cv::fillPoly(roi, contours, COLOR_WEAK);
        }


        //// 4. 画图
        //cv::Mat roi = mFrame(cv::Rect(x1, y1, x2 - x1, y2 - y1));
        ////cv::Scalar color = (mask_area / hole_area) < AREA_THERSHOLD ? COLOR_WEAK : COLOR_EXITS;
        //std::vector<std::vector<cv::Point>> contours;
        //cv::Mat hierarchy;
        //cv::findContours(binary_mask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE); // 查找轮廓
        //cv::fillPoly(roi, contours, COLOR_EXITS);

    }
    
    // 5. 匹配
    std::vector<bool> holes_match_bool;
    hole_target_match(res_bboxes, holes_match_bool);


    // 6. 画出没有预测的边框
    for (int i = 0; i < holes_match_bool.size(); i++) {
        bool is_match = holes_match_bool[i];
        if (!is_match) {
            drawBox(mFrame, hole_xyxys[i]);
        }

        res.push_back(is_match ? 1 : 0);   // 整理结果
    }




}


void model::hole_target_match(const std::vector<std::vector<float>>& target_xyxys, std::vector<bool>& holes_match_bool) {

    /* for (auto box : target_xyxys) {
         drawBox(mFrame, box);
     }
     for (auto box : hole_xyxys) {
         drawBox(mFrame, box);
     }
     cv::imwrite("result.jpg", mFrame);
     exit(1);*/



    std::vector<bool> holes_match(hole_xyxys.size(), false);
    std::vector<bool> targets_match(target_xyxys.size(), false);
    std::vector<std::vector<int>> hole_target_match_vector;

    // ******************4. 一阶段匹配  为每个穴孔寻找匹配目标******************
    for (int hole_id = 0; hole_id < hole_xyxys.size(); hole_id++) {
        auto hole_xyxy_ = hole_xyxys[hole_id];
        int max_id = -1;  //  当前最大的目标id
        float max_score = -1;  //  当前最大的目标分数

        for (int target_id = 0; target_id < target_xyxys.size(); target_id++)
        {
            bool target_match = targets_match[target_id];  // 是否已经匹配
            if (target_match)
            {
                continue;
            }

            auto target_xyxy = target_xyxys[target_id];

            float score_iou = iou(hole_xyxy_, target_xyxy);  // 计算IoU


            // 提前结束循环
            if (score_iou > 0.8)
            {
                max_id = target_id;
                max_score = score_iou;
                break;
            }
            else if (score_iou > max_score)
            {
                max_id = target_id;
                max_score = score_iou;
            }

        }



        if ((max_id != -1) && (max_score >= MATCH_THERSHOLD))
        {
            targets_match[max_id] = true;
            holes_match[hole_id] = true;
            hole_target_match_vector.emplace_back(std::vector<int>{{hole_id, max_id}});

        }


    }




    // ******************4. 二阶段匹配 为剩余的目标 分配穴孔******************
    for (int target_id = 0; target_id < target_xyxys.size(); target_id++)
    {
        bool target_match = targets_match[target_id];  // 是否已经匹配
        if (target_match)
        {
            continue;
        }

        auto target_xyxy = target_xyxys[target_id];
        int max_id = -1;  //  当前最大的目标id
        float max_score = -1;  //  当前最大的目标分数


        for (int hole_id = 0; hole_id < hole_xyxys.size(); hole_id++)
        {
            /*if (holes_match[hole_id])
            {
                continue;
            }*/
            auto hole_xyxy_ = hole_xyxys[hole_id];
            float score_iou = iou(hole_xyxy_, target_xyxy);  // 计算IoU

            // 提前结束循环
            if (score_iou > 0.8)
            {
                max_id = hole_id;
                max_score = score_iou;
                break;
            }
            else if (score_iou > max_score)
            {
                max_id = hole_id;
                max_score = score_iou;
            }
        }

        if ((max_id != -1) && max_score >= (MATCH_THERSHOLD * 0.2))
        {
            targets_match[target_id] = true;
            holes_match[max_id] = true;
            hole_target_match_vector.emplace_back(std::vector<int>{{max_id, target_id}});

        }
    }

    holes_match_bool = holes_match;

};


std::vector<int> model::infer(cv::Mat img) {
    // 1. 数据预处理，获得输入数据，std::vector<float>
    preProcess(img); // 前处理
    // 2. 前向推理，得到结果std::vector<float>
    forward();
    // 3. 后处理
    std::vector<int> res;
    postProcess(res);
    return res; // 返回
};





void model::save_img(const std::string save_path) {
    cv::imwrite(save_path, mFrame);
}


void drawBox(cv::Mat& img, std::vector<float>& b_xyxyc) {

    // 获取边界框的坐标
    int x1 = static_cast<int>(b_xyxyc[0]);
    int y1 = static_cast<int>(b_xyxyc[1]);
    int x2 = static_cast<int>(b_xyxyc[2]);
    int y2 = static_cast<int>(b_xyxyc[3]);
    cv::rectangle(img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(0, 0, 255), 2);
};


bool get_hole_xyxys(std::vector<std::vector<float>>& hole_xyxys, int w, int h) {
    // 计算每个单元格的宽度和高度
    float cell_width = w / HOLE_NUM_W;
    float cell_height = h / HOLE_NUM_H;
    // 存储掩码的容器

    hole_xyxys.clear();
    // 遍历每个单元格
    for (int i = 0; i < HOLE_NUM_H; ++i) { // 纵向
        float y1 = static_cast<float>(i * cell_height);
        float y2 = static_cast<float>((i + 1) * cell_height);
        for (int j = 0; j < HOLE_NUM_W; ++j) { // 横向
            float x1 = static_cast<float>(j * cell_width);
            float x2 = static_cast<float>((j + 1) * cell_width);
            // 创建掩码矩阵
            std::vector<float> xyxy = { x1, y1, x2, y2 };  // 单通道，8位无符号整数
            hole_xyxys.push_back(xyxy);
        }
    }
}


float iou(const std::vector<float> hole_box1, const std::vector<float> box2) {


    float x1 = MAX(hole_box1[0], box2[0]);
    float y1 = MAX(hole_box1[1], box2[1]);
    float x2 = MIN(hole_box1[2], box2[2]);
    float y2 = MIN(hole_box1[3], box2[3]);

    // 计算重叠面积
    auto width = MAX(0, (x2 - x1));
    auto height = MAX(0, (y2 - y1));
    float inter_area = (float)(width * height);


    float hole_area = (float)(hole_box1[2] - hole_box1[0]) * (hole_box1[3] - hole_box1[1]);
    float iou_v = inter_area / (hole_area + 1e-9);


    return iou_v;

};


//  202507-07-14